import tensorflow_hub as hub
import numpy as np 

class BigTransfer:

    def __init__(self):
        # big Transfer D:\project\Image - Similarity\pretain_model
        # efficientnet E:\nlp\pretrain_models\image\efficientnet_b7_feature-vector_1

        use_type = "efficientnet_v2"

        info = {
            "efficientnet_v2":{
                "url":r"E:\nlp\pretrain_models\image\tf-hub\imagenet_efficientnet_v2_imagenet21k_xl_feature_vector_2",
                "size":512
            },
            "efficientnet_v2_ft":{
                'url':r'E:\nlp\pretrain_models\image\tf-hub\imagenet21k-ft1k-xl-feature-vector',
                "size":512
            }
            ,"inception_v3":{
                'url':r"E:\nlp\pretrain_models\image\tf-hub\inception_v3",
                "size": 299
            }
        }

        self.module = hub.KerasLayer(info[use_type]["url"])

        self.module.build([None, info[use_type]["size"], info[use_type]["size"], 3])

    def extract_feature(self, imgs):
        imgs_batch = np.concatenate(imgs,axis=0)

        features = self.module(imgs_batch)
        return np.array(features).tolist()